# -*- coding: utf-8 -*-
"""
Created on Sun Oct 11 10:09:35 2020
用于库伦耦合的半导体双量子点模型
128初始态对制备保真度

@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
np.random.seed(1)
I = np.matrix(np.identity(2, dtype=complex))
sx = np.mat([[0, 1], [1, 0]], dtype=complex)
sy = np.mat([[0, -1j], [1j, 0]], dtype=complex)
sz = np.mat([[1, 0], [0, -1]], dtype=complex)
h_1 = 1
h_2 = 1
coupling = 1/2
# init_psi = np.mat([[1], [0], [0], [0]], dtype=complex)
# target_psi = np.mat([[1], [0], [0], [1]], dtype=complex)/np.sqrt(2) #最终目标态
init_set_num = 512
target_set_num = init_set_num

action_space =np.array([[1,1],
                        [1,2],
                        [1,3],
                        [1,4],
                        [1,5],
                        [2,1],
                        [2,2],
                        [2,3],
                        [2,4],
                        [2,5],
                        [3,1],
                        [3,2],
                        [3,3],
                        [3,4],
                        [3,5],
                        [4,1],
                        [4,2],
                        [4,3],
                        [4,4],
                        [4,5],
                        [5,1],
                        [5,2],
                        [5,3],
                        [5,4],
                        [5,5]
                        ])
n_actions = len(action_space)

alpha_num = 4
theta = [np.pi/8,np.pi/4,3*np.pi/8]
theta_1 = theta
theta_2 = theta
theta_3 = theta

alpha = np.linspace(0,np.pi*2,alpha_num,endpoint=False)
alpha_1 = alpha
alpha_2 = alpha
alpha_3 = alpha
alpha_4 = alpha

a_list_complex = np.matrix([[0,0,0,0]],dtype=complex) #第一行用来占位，否则无法和其他行并在一起，在最后要注意去掉这一行
for ii in range(3): #theta_1
    for jj in range(3): #theta_2
        for kk in range(3): #theta_3
            for mm in range(alpha_num): #alpha_1
                for nn in range(alpha_num): #alpha_2
                    for oo in range(alpha_num): #alpha_3
                        for pp in range(alpha_num): #alpha_4
                            
                            a_1_mo = np.cos(theta_1[ii])
                            a_2_mo = np.sin(theta_1[ii])*np.cos(theta_2[jj])
                            a_3_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.cos(theta_3[kk])
                            a_4_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.sin(theta_3[kk])
                            
                            a_1_real = a_1_mo*np.cos(alpha_1[mm])
                            a_1_imag = a_1_mo*np.sin(alpha_1[mm])
                            a_2_real = a_2_mo*np.cos(alpha_2[nn])
                            a_2_imag = a_2_mo*np.sin(alpha_2[nn])
                            a_3_real = a_3_mo*np.cos(alpha_3[oo])
                            a_3_imag = a_3_mo*np.sin(alpha_3[oo])
                            a_4_real = a_4_mo*np.cos(alpha_4[pp])
                            a_4_imag = a_4_mo*np.sin(alpha_4[pp])
                            
                            a_1_complex = a_1_real + a_1_imag*1j
                            a_2_complex = a_2_real + a_2_imag*1j
                            a_3_complex = a_3_real + a_3_imag*1j
                            a_4_complex = a_4_real + a_4_imag*1j
                            
                            a_complex = np.matrix([[ a_1_complex, a_2_complex, a_3_complex, a_4_complex]])
                            a_list_complex = np.row_stack((a_list_complex,a_complex))
                            
psi_set = np.array(np.delete(a_list_complex,0,axis=0)) # 删除矩阵的第一行
np.random.shuffle(psi_set) #打乱顺序
init_set = psi_set[:init_set_num]

def target_set():
    target_set = psi_set[init_set_num : init_set_num + target_set_num]
    return target_set

#动作直接选最优的
def step(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
        J_12 = J_1 * J_2 /2
    
        H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                        J_12/2*np.kron((sz - I),(sz - I)) + \
           h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_


#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
        J_12 = J_1 * J_2 /2
    
        H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                        J_12/2*np.kron((sz - I),(sz - I)) + \
           h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling

        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
    
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        # del action_list[fid_list.index(max(fid_list))]
        # del psi_list[fid_list.index(max(fid_list))]
        # del fid_list[fid_list.index(max(fid_list))]
        
        # best_action = fid_list.index(max(fid_list))
        # best_fid = max(fid_list)
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
        J_12 = J_1 * J_2 /2
    
        H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                        J_12/2*np.kron((sz - I),(sz - I)) + \
           h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        del action_list[fid_list.index(max(fid_list))]
        del psi_list[fid_list.index(max(fid_list))]
        del fid_list[fid_list.index(max(fid_list))]
        
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
        
        # best_action = fid_list.index(min(fid_list))
        # best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_
#---------------------------------------------------------------------------------
#将数值列表从小到大进行排序（冒泡法）
def sort_fid_list(lt):
    n= len(lt)
    for x in range(n-1):
       for y in range(n-1-x):
          if lt[y]>lt[y+1]:
             lt[y],lt[y+1]=lt[y+1],lt[y]
    return lt
#--------------------------------------------------------------------------------




time1 = time()
def job(target_psi):
    ffid_list = np.zeros((init_set_num ,1))
    target_psi = np.mat(target_psi).T
    for k in range(init_set_num):
        psi1 = np.mat(init_set[k]).T
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        fid_max = 0
        fid_max1 = 0
        fid_max2 = 0
        fid_max0 = 0
        step_n = 0
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
            
        fid_max = max(fid_max1,fid_max2)
        if fid_max < 0.999:
            step_n = 0
            F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
            psi = psi1
            while True:
                action, F, psi_ = step(psi,target_psi,F)
                fid_max0 = max(F,fid_max0)
                psi = psi_
                step_n += 1
                if fid_max0>0.999 or step_n>step_max:
                    break 
            fid_max = max(fid_max,fid_max0)  
        ffid_list[k] = fid_max
    return  np.mean(ffid_list)


def multicore():
    pool = mp.Pool()
    F_list = pool.map(job, target_set)
    return F_list
    
if __name__ == '__main__':
    target_set = target_set()
    target_set = [target_set[48]]
    # target_set = [init_set[48]]  # the worst point 
    time1 = time()
    print('running...')
    
    for kT in [5,8,10,12,15,18,20]:#:
        for kdt in [2,3,4,5,6,8,10]:#
            F_list = []
            T = kT*np.pi
            dt = np.pi/kdt
            step_max = T/dt
        
            F_list = multicore()
            # print(F_list)
            print("kT = ",kT,"kdt = ",kdt,"mean_fid = ",np.mean(F_list))
            # print(np.mean(F_list))
    time2 = time()
    print('time_cost is: ',time2-time1)
    
# running...
# kT =  5 kdt =  2 mean_fid =  0.7438358088170904
# kT =  5 kdt =  3 mean_fid =  0.6897224809311756
# kT =  5 kdt =  4 mean_fid =  0.6771562100763209
# kT =  5 kdt =  5 mean_fid =  0.6787576074819761
# kT =  5 kdt =  6 mean_fid =  0.7077164224659412
# kT =  5 kdt =  8 mean_fid =  0.7333985552502822
# kT =  5 kdt =  10 mean_fid =  0.7199544974591816
# kT =  8 kdt =  2 mean_fid =  0.7663399292206176
# kT =  8 kdt =  3 mean_fid =  0.718433756526815
# kT =  8 kdt =  4 mean_fid =  0.710486487539852
# kT =  8 kdt =  5 mean_fid =  0.700015478296118
# kT =  8 kdt =  6 mean_fid =  0.7296991028825691
# kT =  8 kdt =  8 mean_fid =  0.7512162991656108
# kT =  8 kdt =  10 mean_fid =  0.7399972061802693
# kT =  10 kdt =  2 mean_fid =  0.7793023335881848
# kT =  10 kdt =  3 mean_fid =  0.7304865724153021
# kT =  10 kdt =  4 mean_fid =  0.7195613902756187
# kT =  10 kdt =  5 mean_fid =  0.7120101449967371
# kT =  10 kdt =  6 mean_fid =  0.7401546889705284
# kT =  10 kdt =  8 mean_fid =  0.758960891747402
# kT =  10 kdt =  10 mean_fid =  0.7457133826723026
# kT =  12 kdt =  2 mean_fid =  0.7873777990482018
# kT =  12 kdt =  3 mean_fid =  0.7410737987785879
# kT =  12 kdt =  4 mean_fid =  0.7294297557184176
# kT =  12 kdt =  5 mean_fid =  0.7193661453449531
# kT =  12 kdt =  6 mean_fid =  0.7464224373205294
# kT =  12 kdt =  8 mean_fid =  0.7642305641254166
# kT =  12 kdt =  10 mean_fid =  0.7506654108041875
# kT =  15 kdt =  2 mean_fid =  0.7962591896671471
# kT =  15 kdt =  3 mean_fid =  0.7495246613458448
# kT =  15 kdt =  4 mean_fid =  0.7390233943718851
# kT =  15 kdt =  5 mean_fid =  0.7296014207298585
# kT =  15 kdt =  6 mean_fid =  0.7535810604273495
# kT =  15 kdt =  8 mean_fid =  0.7709758380986815
# kT =  15 kdt =  10 mean_fid =  0.7561132219306897
# kT =  18 kdt =  2 mean_fid =  0.8048629956639765
# kT =  18 kdt =  3 mean_fid =  0.7589336470793201
# kT =  18 kdt =  4 mean_fid =  0.7464646821238139
# kT =  18 kdt =  5 mean_fid =  0.7364996878098748
# kT =  18 kdt =  6 mean_fid =  0.7599437359445365
# kT =  18 kdt =  8 mean_fid =  0.7749444580882513
# kT =  18 kdt =  10 mean_fid =  0.7586433923147613
# kT =  20 kdt =  2 mean_fid =  0.8089513262263461
# kT =  20 kdt =  3 mean_fid =  0.7634182660762836
# kT =  20 kdt =  4 mean_fid =  0.75285766486562
# kT =  20 kdt =  5 mean_fid =  0.7422638753775486
# kT =  20 kdt =  6 mean_fid =  0.7635345188092106
# kT =  20 kdt =  8 mean_fid =  0.7774900701928051
# kT =  20 kdt =  10 mean_fid =  0.7604147958403429
# time_cost is:  54832.24940896034

    
